import json
import re
from enum import Enum
from typing import List

from pydantic import BaseModel, Field, field_validator


class Scene(BaseModel):
    label: str = Field(description="场景引导关键词")
    category: str = Field("", description="场景引导类别")

    @classmethod
    def fix_str(cls, value: str, scene2category: dict) -> "Scene":
        if not value or not isinstance(value, str):
            return ""
        items = re.split(r"[,，、]", value.strip())
        results = []
        for item in items:
            if "&" not in item:
                if item not in scene2category:
                    continue
                else:
                    results.append(f"{item}&{scene2category[item]}")
            else:
                results.append(item)

        return "，".join(results)

    def to_str(self) -> str:
        return f"{self.category}&{self.label}"


class Graph(BaseModel):
    level_1: List[str] = Field(description="一级事项标签")
    level_2: List[str] = Field(description="二级事项标签")
    scene: List[Scene] = Field(description="场景引导标签列表")
    keywords: List[str] = Field(default=[], description="关键词列表")

    @field_validator("level_1", mode="before")
    def transform_level_1(cls, value):
        if not isinstance(value, str) and not isinstance(value, list):
            return []
        elif isinstance(value, str):
            return re.split(r"[,，]", value.strip())
        else:
            return value

    @field_validator("level_1", mode="after")
    def validate_empty_list(cls, value):
        return [item for item in value if item]

    @field_validator("level_2", mode="before")
    def transform_level_2(cls, value):
        if not isinstance(value, str) and not isinstance(value, list):
            return []
        elif isinstance(value, str):
            return re.split(r"[,，]", value.strip())
        else:
            return value

    @field_validator("level_2", mode="after")
    def validate_level_2(cls, value):
        return [item for item in value if item]

    @field_validator("scene", mode="before")
    def transform_scene(cls, value):
        if not isinstance(value, str) and not isinstance(value, list) or not value:
            return []
        elif isinstance(value, str):
            results = []
            scene_list = re.split(r"[,，、]", value.strip())

            for scene_str in scene_list:
                if "&" not in scene_str:
                    raise ValueError(
                        f"Invalid scene format: {scene_str}. Scene should be in the format of 'label&category'."
                    )
                items = scene_str.split("&")
                if len(items) != 2:
                    raise ValueError(
                        f"Invalid scene format: {scene_str}. Scene should be in the format of 'label&category'."
                    )

                category, labels = items

                for label in labels.split("/"):
                    results.append(Scene(label=label, category=category).model_dump())

            return results
        else:
            return value

    @field_validator("scene", mode="after")
    def validate_scene(cls, value):
        return [item for item in value if item.label]

    @field_validator("keywords", mode="before")
    def transform_keywords(cls, value):
        if not isinstance(value, str) and not isinstance(value, list):
            return []
        elif isinstance(value, str):
            return re.split(r"[,，、]", value.strip())
        else:
            return value

    @field_validator("keywords", mode="after")
    def validate_keywords(cls, value):
        return [item for item in value if item]

    def to_str(self) -> str:
        """Convert the graph to a string."""
        tags = []
        tags.extend(self.level_1)
        tags.extend(self.level_2)
        tags.extend([scene.label for scene in self.scene])
        tags.extend(self.keywords)

        return "要办理的事项：" + "，".join(tags)

    @classmethod
    def from_str(cls, graph_str: str) -> "Graph":
        """Create a graph from a string."""
        level_1 = []
        level_2 = []
        scene = []
        keywords = []
        for line in graph_str.split("\n"):
            if line.startswith("一级事项："):
                level_1 = line.replace("一级事项：", "").split(", ")
            elif line.startswith("二级事项："):
                level_2 = line.replace("二级事项：", "").split(", ")
            elif line.startswith("场景引导："):
                scene = [
                    Scene(
                        label=scene_str,
                        category="",
                    ).model_dump()
                    for scene_str in line.replace("场景引导：", "").split(", ")
                ]
            elif line.startswith("关键词："):
                keywords = line.replace("关键词：", "").split(", ")
        return cls(level_1=level_1, level_2=level_2, scene=scene, keywords=keywords)

    @classmethod
    def from_llm_response(cls, response: str) -> "Graph":
        # 提取 JSON 部分
        json_match = re.findall(r"^\{.*^\}", response, re.DOTALL | re.MULTILINE)
        if json_match:
            json_string = json_match[-1]
            # 替换伪 JSON 为标准 JSON
            json_string = json_string.replace("'", '"').replace(": ", ":")
            try:
                parsed_output = json.loads(json_string)
                return cls.model_validate(parsed_output)
            except json.JSONDecodeError:
                return cls(level_1=[], level_2=[], scene=[], keywords=[])
        else:
            return cls(level_1=[], level_2=[], scene=[], keywords=[])

    def update(self, other: "Graph", kg_tree: dict) -> None:
        """Merge the graph with another graph."""

        if not set(other.level_1).issubset(set(self.level_1)):
            self.level_1 = other.level_1
            # remove old level_2 not in new level_1
            remove_list = []
            for i, ol2 in enumerate(self.level_2):
                if ol2 not in kg_tree.get(self.level_1[0], {}):
                    remove_list.append(i)
            for i in remove_list[::-1]:
                self.level_2.pop(i)

        if other.level_2:
            self.level_2 = other.level_2

        # Merge scenes with same category
        category_scenes = {}
        # First add existing scenes
        for scene in self.scene:
            if scene.category not in category_scenes:
                category_scenes[scene.category] = set()
            category_scenes[scene.category].add(scene.label)

        # Then merge new scenes
        for scene in other.scene:
            if scene.category not in category_scenes:
                category_scenes[scene.category] = set()
            category_scenes[scene.category].add(scene.label)

        # Convert back to Scene objects
        self.scene = [
            Scene(label=label, category=category)
            for category, labels in category_scenes.items()
            for label in labels
        ]

        self.keywords = other.keywords

    def is_empty(self) -> bool:
        """Check if the graph is empty."""
        return not any([self.level_1, self.level_2, self.scene])

    def deduplicate(self) -> None:
        """Deduplicate the graph fields."""
        self.level_1 = list(set(self.level_1))
        self.level_2 = list(set(self.level_2))
        self.keywords = list(set(self.keywords))

    def match_score(self, other: "Graph") -> float:
        """Calculate the match score between two graphs."""
        score = 0
        if self.level_1 and other.level_1:
            score += len(set(self.level_1).intersection(other.level_1))
        if self.level_2 and other.level_2:
            score += len(set(self.level_2).intersection(other.level_2))
        if self.scene and other.scene:
            score += len(
                set([scene.label for scene in self.scene]).intersection(
                    [scene.label for scene in other.scene]
                )
            )
        score = score / (len(self.level_1) + len(self.level_2) + len(self.scene))
        return score


class LabelType(str, Enum):
    """标签类型枚举"""

    LEVEL_1 = "level_1"
    LEVEL_2 = "level_2"
    SCENE = "scene"
    KEYWORDS = "keywords"


class ALL_NODES_TYPE(BaseModel):
    all_level_1: List[str]
    all_level_2: List[str]
    all_scene: dict[str, List[str]]
    all_keywords: List[str]
    all_scene_category: List[str]
    all_scene_label: List[str]
    scene_label_to_category: dict[str, str]


def refine_graph(kg_tree: dict, graph: Graph, all_nodes: ALL_NODES_TYPE) -> Graph:
    """Refine the graph by using the knowledge graph tree.

    Args:
        kg_tree: The knowledge graph tree structure mapping level_1 -> level_2 -> category -> scenes
        graph: The graph to refine
    """
    graph = graph.model_copy()

    # 这个逻辑暂时移除
    # # 如果没有抽取到二级事项，一级事项下面又只有一个二级事项，那么将这个二级事项加入到二级事项列表中
    # if graph.level_1:
    #     for level_1 in graph.level_1:
    #         if level_1 not in kg_tree:
    #             continue
    #         if len(kg_tree[level_1]) == 1:
    #             graph.level_2.extend(kg_tree[level_1].keys())

    # 众神归位，修复标签级别错误的问题
    # TODO 场景引导没处理
    for label in graph.level_1:
        if label not in all_nodes.all_level_1 and label in all_nodes.all_level_2:
            graph.level_1.remove(label)
            graph.level_2.append(label)

    for label in graph.level_2:
        if label not in all_nodes.all_level_2 and label in all_nodes.all_level_1:
            graph.level_2.remove(label)
            graph.level_1.append(label)

        elif label not in all_nodes.all_level_2 and label in all_nodes.scene_label_to_category:
            graph.level_2.remove(label)
            graph.scene.append(
                Scene(label=label, category=all_nodes.scene_label_to_category[label])
            )
        elif label not in all_nodes.all_level_2:
            graph.level_2.remove(label)
            graph.keywords.append(label)

    remove_list = []
    for i, scene in enumerate(graph.scene):
        # 防止位置出错
        if (
            scene.category in all_nodes.all_scene_label
            and scene.label in all_nodes.all_scene_category
        ):
            # swap scene.category and scene.label
            scene.category, scene.label = scene.label, scene.category
        elif scene.category not in all_nodes.all_scene_category:
            remove_list.append(i)
        elif scene.label not in all_nodes.all_scene_label:
            remove_list.append(i)
            if scene.label in all_nodes.all_level_2:
                graph.level_2.append(scene.label)
            elif scene.label in all_nodes.all_level_1:
                graph.level_1.append(scene.label)

    graph.scene = [scene for i, scene in enumerate(graph.scene) if i not in remove_list]

    # TODO 社保专门处理
    if "社保" in graph.level_1 and len(graph.level_2) > 0:
        cur_level_2 = kg_tree["社保"].keys()
        flag = False
        for level_2 in cur_level_2:
            if level_2 in graph.level_2:
                flag = True
                break
        if not flag:
            graph.level_1.remove("社保")

    # 如果一级事项为空，而且二级事项不为空，那么根据二级事项找到对应的一级事项(如果只有一个)
    existing_level_1 = set(graph.level_1)
    # Find missing level_1 items by checking each level_2 item
    for level_2 in graph.level_2:
        cur_level_1 = set()
        # Find all level_1 items that contain this level_2
        for level_1, level_2_items in kg_tree.items():
            if level_2 in level_2_items and level_1 not in existing_level_1:
                existing_level_1.add(level_1)
                cur_level_1.add(level_1)

        if len(cur_level_1) == 1 and not graph.level_1:
            graph.level_1.append(cur_level_1.pop())

    graph.deduplicate()
    return graph